from maya import cmds
from maya.api import OpenMaya as OM2
from maya.api import OpenMayaAnim as OMA2

import baseUtils, shapeUtils, geoUtils, attrUtils
import utils, namespaceUtils, fileManage, mathNodeCalculate, weightCalculate

current_path = fileManage.current_maya_path()


def is_blendshape(blendshape):
    """Check if the node is a valid blendshape node or not.

    Args:
        blendshape (str): Name of the blendshape to check.

    Returns:
        bool: True if the node is a valid blendshape node or False if not.

    """

    # Check blendShape exists
    if not cmds.objExists(blendshape) or cmds.objectType(blendshape) != 'blendShape':

        return False

    return True


def get_blend_shapes(node):
    """Get all the blendshapes nodes on the node.

    Args:
        node (str): The node to query all the blendshapes connected to it.

    Returns:
        list/None: A list of all the blendshapes nodes on the node.

    """

    shape = shapeUtils.get_shape(node=node)
    if shape:
        history = cmds.listHistory(shape, pruneDagObjects=True, il=2)
        if history:
            blendshape_list = [x for x in history if is_blendshape(blendshape=x)]

            return blendshape_list

        return None

    cmds.warning('No shape is connected to {}.'.format(node))

    return None


def create_blendshape(base_geo, name=None, origin='local', deform_order=None):
    """Create a blendShape deformer for the specified geometry.

    Args:
        base_geo (str): Geometry to apply blendShape deformer to.
        name (str): Blendshape name.
        origin (str): Create a local or world space blendShape deformer. 'local'/'world'.
        deform_order (str/None): Deformer order. 'after'/'before'/'parallel'/'split'/'foc'/default if None.

    Returns:
        str: The newly created blendshape name or the current name if it already exists.

    """

    # If name is not specified
    if not name:
        name = '{}_blendShape'.format(namespaceUtils.remove_namespace(node=base_geo))
        # create new name with number after if the name is already exist
        all_nodes = cmds.ls()
        highest_num = fileManage.find_highest_trailing_number(names=all_nodes, base_name=name)
        name = '{}{}'.format(name, highest_num+1)

    # Check BlendShape
    if is_blendshape(blendshape=name):
        # print('BlendShape {} already exist!'.format(name))
        return name

    # Create BlendShape
    if deform_order == 'after':
        blendshape = cmds.blendShape(base_geo, name=name, origin=origin, after=True)[0]
    elif deform_order == 'before':
        blendshape = cmds.blendShape(base_geo, name=name, origin=origin, before=True)[0]
    elif deform_order == 'parallel':
        blendshape = cmds.blendShape(base_geo, name=name, origin=origin, parallel=True)[0]
    elif deform_order == 'split':
        blendshape = cmds.blendShape(base_geo, name=name, origin=origin, split=True)[0]
    elif deform_order == 'foc':
        blendshape = cmds.blendShape(base_geo, name=name, origin=origin, foc=True)[0]
    else:
        blendshape = cmds.blendShape(base_geo, name=name, origin=origin)[0]

    return blendshape


class BlendShapes:
    def __init__(self, base_geo, blendshape_name):
        """Create blendshape node, and access all the blendshape functions

        Args:
            base_geo (str): Geometry to apply blendShape deformer to.
            blendshape_name (str): Blendshape name.

        Raises:
            IndexError: If no base geometry is specified.
            ValueError: If base geometry doesn't have shape node.
            Exception: If Unable to determine deformer set for this blendshape.

        """

        # Get the base geometry
        if not base_geo:
            try:
                base_geo = cmds.ls(sl=True)[0]
            except (IndexError, Exception):
                raise IndexError('No base geometry is selected!')

        # Check if base geometry has shape node
        self.base_geo_shape = shapeUtils.get_shape(node=base_geo)
        if not self.base_geo_shape:
            raise ValueError('No shape connected to {}!'.format(base_geo))

        # Define base geometry
        self.base_geo = base_geo
        # Get base geometry object type, mesh, nurbsCurve, nurbsSurface, lattice.
        self.base_geo_type = cmds.objectType(self.base_geo_shape)
        # Shape M Object
        self.shape_m_obj = baseUtils.get_dag_obj(obj=self.base_geo_shape)

        # Get blendshape node
        self.blendshape = create_blendshape(base_geo=self.base_geo, name=blendshape_name, origin='local')

        # Define blendshape json file name
        self.blendshape_data_file = '{}_{}_data.json'.format(
            namespaceUtils.remove_namespace(node=self.base_geo), self.blendshape
        )

        self.targets_grp = '{}_targets_Grp'.format(self.blendshape)

        self.inbetween_targets_grp = '{}_inbetween_targets_Grp'.format(self.blendshape)
        self.extracted_geos_grp = '{}_extracted_geos_Grp'.format(self.blendshape)

        # Setting groups
        for grp in [self.targets_grp, self.extracted_geos_grp, self.inbetween_targets_grp]:
            if not cmds.objExists(grp):
                cmds.group(empty=True, name=grp)
            cmds.setAttr('{}.v'.format(grp), 0)

        # Get deformer fn (blendshape)
        self.blendshape_fn = self._get_blendshape_fn()
        # Get base geometry index for inputTarget in this blendshape deformer
        self.base_geo_index = self._get_base_geo_index()


        # Get blendshape targets index dict
        self.targets_index_dict = {}
        # Get blendshape targets attributes weights dict
        self.targets_attributes_weights = {}
        # Get blendshape targets inbetween dict
        self.targets_inbetween_dict = attrUtils.get_string_info(obj=self.blendshape, attr='Targets_Inbetween_Dict')
        if not self.targets_inbetween_dict:
            self.targets_inbetween_dict = {}
        # Get blendshape targets weights dict
        self.targets_weights_dict = {}
        # Get weight to blendshape targets connection dict
        self.targets_connections_dict = attrUtils.get_string_info(obj=self.blendshape, attr='Targets_Connections_Dict')
        if not self.targets_connections_dict:
            self.targets_connections_dict = {}

        # get the targets groups information list. The list of hierarchies dictionaries
        # with key ('parent'), value (group name), key (children), value (targets names)
        self.targets_groups_list = attrUtils.get_string_info(obj=self.blendshape, attr='Targets_Groups_List')
        if not self.targets_groups_list:
            self.targets_groups_list = [{'parent': '', 'children': []}]

        # Default split locators
        self.pos_loc, self.neg_loc = None, None
        split_locators = attrUtils.get_string_info(obj=self.blendshape, attr='Split_Locators')
        if split_locators:
            self.pos_loc, self.neg_loc = split_locators



        # Get blendshape set

        # 尝试从 API 获取 deform set
        try:
            self.blendshape_set = self.blendshape_fn.deformerSet()
        except RuntimeError:
            self.blendshape_set = None

        # 如果失败，用 cmds fallback 获取 objectSet 名称
        if not self.blendshape_set or self.blendshape_set.isNull():
            set_name = self._get_deformer_set_cmds()
            if set_name:
                sel = OM2.MSelectionList()
                sel.add(set_name)
                self.blendshape_set = sel.getDependNode(0)
            else:
                # 警告但不 raise，可以选择性忽略
                OM2.MGlobal.displayWarning(
                    '[BlendShapes] 无法找到 {} 的 deformerSet，将跳过 set 相关功能。'.format(self.blendshape))
                self.blendshape_set = None
                self.blendshape_set_name = None
                self.members = None
                self.members_components = []
                self.members_indices = []
                return  # 或者你可以选择继续初始化其他部分


        # Get the function set for this set
        self.blendshape_fn_set = OM2.MFnSet(self.blendshape_set)
        self.members = self.blendshape_fn_set.getMembers(True)
        # Get components as a list
        self.members_components = self._get_members_list()
        # Get components indices as a list
        self.members_indices = self._get_members_indices()




        # Define blendshape date file to save into json
        self.blendshape_data = {
            'name': self.blendshape,
            'geometry': self.base_geo,
            'shapeType': self.base_geo_type,
            'members': self.members_components,
            'members_indices': self.members_indices,
            'targets_index_dict': self.targets_index_dict,
            'targets_inbetween_dict': self.targets_inbetween_dict,
            'targets_weights_dict': self.targets_weights_dict,
            'targets_connections_dict': self.targets_connections_dict,
            'targets_groups_list': self.targets_groups_list
        }


    def _get_deformer_set_cmds(self):
        """Try to find the deformer set via cmds as a fallback."""
        sets = cmds.listConnections(self.blendshape + ".message", type="objectSet") or []
        for s in sets:
            if cmds.sets(s, q=True):  # ensure it's a real set with members
                return s
        return None


    def _get_blendshape_fn(self):
        """Initializes a new MFnGeometryFilter function set and attaches it to a geometryFilter node.

        Returns:
            OpenMayaAnim.MFnGeometryFilter: The Function set for operating on geometryFilter nodes.

        """

        return OMA2.MFnGeometryFilter(baseUtils.get_dag_obj(obj=self.blendshape))

    def _get_base_geo_index(self):
        """Returns the plug index for the specified output shape.

        Returns:
            long: The plug index for the specified output shape.

        """

        return self.blendshape_fn.indexForOutputShape(self.shape_m_obj)

    def rename_blendshape_set(self, set_name=''):
        """Rename the blendshape deformer set name.

        Args:
            set_name (str): New set name.

        """

        if not set_name:
            set_name = '{}_Set'.format(self.blendshape)

        if self.blendshape_set_name != set_name:
            cmds.rename(self.blendshape_set_name, set_name)


    def _get_members_list(self):
        """Returns deformed components as a list for the deformed geometry.

        Returns:
            list: The components as a list for the deformed geometry.

        """

        members_list = self.members.getSelectionStrings()

        return cmds.ls(members_list[self.base_geo_index], flatten=True)

    def _get_members_indices(self):
        """Returns the indexed deformed components as a list for the deformed geometry,
           works on 'mesh'/'nurbsCurve'/nurbsSurface'/'lattice'/.

        Returns:
            list: The indexed components as a list.

        """

        dag_path, components = self._get_geometry_components()

        # Get set member indices
        members_indices_list = []

        # Single index
        if self.base_geo_type == 'mesh':
            single_index_fn_comp = OM2.MFnSingleIndexedComponent(components)
            members_indices_list = list(single_index_fn_comp.getElements())

        elif self.base_geo_type == 'nurbsCurve':
            single_index_fn_comp = OM2.MFnSingleIndexedComponent(components)
            members_indices_list = list(single_index_fn_comp.getElements())

        # Double index
        elif self.base_geo_type == 'nurbsSurface':
            double_index_fn_comp = OM2.MFnDoubleIndexedComponent(components)
            members_indices_list = double_index_fn_comp.getElements()

        # Triple Index
        elif self.base_geo_type == 'lattice':
            triple_index_fn_comp = OM2.MFnTripleIndexedComponent(components)
            members_indices_list = triple_index_fn_comp.getElements()

        return members_indices_list

    def add_targets(self, targets, topology_check=False):
        """Add new targets to the blendshape.

        Args:
            targets (list): The new targets list to add to the blendshape.
            topology_check (bool): Check the topology match between the shapes being blended or not

        Returns:
            list: Newly added targets list.

        """

        # Get the current highest target id, then get newly added target id (after highest)
        highest_index = self._get_highest_index()

        # get targets shapes groups
        targets_group = [self.targets_groups_list[i]['parent'] for i in range(len(self.targets_groups_list))]

        # Only add those targets with shape node.
        targets_to_add_list = [
            target
            for
            target
            in
            targets
            if shapeUtils.get_shape(node=target)
            and cmds.objectType(shapeUtils.get_shape(node=target)) == self.base_geo_type
        ]
        targets_list = []

        for i in range(len(targets_to_add_list)):
            target_to_add = namespaceUtils.remove_namespace(node=targets_to_add_list[i])
            if target_to_add not in self.targets_index_dict and target_to_add not in targets_group:
                cmds.blendShape(
                    self.blendshape,
                    e=True,
                    t=[self.base_geo, highest_index+i+1, targets_to_add_list[i], 1.0],
                    topologyCheck=topology_check)

                self.targets_index_dict[target_to_add] = highest_index+i+1
                # add to targets inbetween dict
                self.targets_inbetween_dict[target_to_add] = {}
                # add to targets groups list
                self.targets_groups_list[0]['children'].append(target_to_add)
                targets_list.append(target_to_add)

            else:
                cmds.warning(
                    '{} target name already exists in this blendshape node as target or group, '
                    'chose another one!'.format(target_to_add)
                )

        attrUtils.add_string_info(
            information=self.targets_inbetween_dict, obj=self.blendshape, attr='Targets_Inbetween_Dict'
        )
        attrUtils.add_string_info(
            information=self.targets_groups_list, obj=self.blendshape, attr='Targets_Groups_List'
        )

        return targets_list

    def add_inbetween_target(self, inbetween_target_geo, target, weight):
        """Add inbetween target of the target in blendshape.

        Args:
            inbetween_target_geo (str): The geometry which needs to be added as inbetween target
                                        to the target in blendshape.
            target (str): The blendshape target needs to have an inbetween target.
            weight (float): The weight at which the inbetween target will be added.

        Raises:
            ValueError: If inbetween target geometry to add doesn't have shape node.
            ValueError: If inbetween target geometry doesn't have same shape type as base geometry.
            RuntimeError: If target already has inbetween target shape at given weight
            RuntimeError: If target doesn't exist in this blendshape node.

        """

        inbetween_target_geo_shape = shapeUtils.get_shape(node=inbetween_target_geo)
        if not inbetween_target_geo_shape:
            raise ValueError('No shape connected to {}!'.format(inbetween_target_geo))

        inbetween_target_geo_type = cmds.objectType(inbetween_target_geo_shape)
        if inbetween_target_geo_type != self.base_geo_type:
            raise ValueError("{} doesn't have same shape type as {}!".format(inbetween_target_geo, self.base_geo))

        target_index = self.get_target_index(target=target)
        if target_index >= 0:
            if weight not in self.targets_inbetween_dict[target]:

                inbetween_target_to_add = namespaceUtils.remove_namespace(node=inbetween_target_geo)
                # check inbetween_target_geo names
                inbetween_target_geo_exist = False
                for weight_inbetween in self.targets_inbetween_dict.values():
                    if inbetween_target_to_add in weight_inbetween.values():
                        inbetween_target_geo_exist = True
                        break

                if not inbetween_target_geo_exist and not cmds.objExists(
                        '{}.{}'.format(self.blendshape, inbetween_target_to_add)
                ):
                    # Set inbetween value if 0< weight < 1
                    if 0.0 < weight < 1.0:
                        cmds.blendShape(
                            self.blendshape,
                            e=True,
                            ib=True,
                            t=[self.base_geo, target_index, inbetween_target_geo, weight],
                            topologyCheck=False
                        )

                        # update to targets inbetween dict
                        self.targets_inbetween_dict[target][weight] = inbetween_target_to_add
                        attrUtils.add_string_info(
                            information=self.targets_inbetween_dict, obj=self.blendshape, attr='Targets_Inbetween_Dict'
                        )

                else:
                    raise RuntimeError(
                        '{} inbetween target already exists as target shape '
                        'or inbetween target shape in this blendshape node'.format(inbetween_target_to_add))

            else:
                raise RuntimeError('{} target already has inbetween target shape at this weight'.format(target))

        else:
            raise RuntimeError('{} target does not exist in this blendshape node!'.format(target))

    def _get_highest_index(self):
        """Return the highest index of the target in this blendshape.

        Returns:
            int: The highest index of the target in this blendshape.

        """

        # Get target index dict
        self.get_targets_index_dict()

        highest_index = 0

        if self.targets_index_dict:
            index_list = list(self.targets_index_dict.values())
            index_list.sort()
            highest_index = index_list[-1]

        return highest_index

    def get_target_index(self, target):
        """Returns the index of the target in this blendshape.

        Args:
            target (str): The target name.

        Returns:
            int/None: The index of the target in this blendshape.

        """

        # Check if target exist
        if cmds.objExists('{}.{}'.format(self.blendshape, target)):
            # Get target index dict
            self.get_targets_index_dict()

            return self.targets_index_dict[target]

        return None
    def get_targets_index_dict(self):
        """Sets the 'self.targets_index_dict' blendshape's targets and indexes dictionary,
           with key(target) and value(target index).

        """

        # Return a list [target0, weight[0], target1, weight[1]....]
        targets_w_index_list = cmds.aliasAttr(self.blendshape, q=True)

        if targets_w_index_list:
            for i in range(0, len(targets_w_index_list), 2):
                # Target name
                target = targets_w_index_list[i]
                # Target weight index, such as 'weight[0]'
                target_w = targets_w_index_list[i + 1]
                # Return the int number 0 from 'weight[0]'
                target_index = int(target_w.split('[')[-1].split(']')[0])

                self.targets_index_dict[target] = target_index


    def set_target_connection(self, expression, target):
        """Connect blendshape target channel with the expression result output.
           expression example: '2*(NodeA.attr+NodeB.attr)*NodeC.attr+NodeD.attr/NodeE.attr'

        Args:
            expression (str): The expression which has its result as the output to blendshape target channel.
            target (str): The name of the target which has its channel to be connected from the output of expression.

        Raises:
            RuntimeError: If target already has an input connection in this blendshape node.
            RuntimeError: If target doesn't exist in this blendshape node.

        """

        # If target name exists and has no connection input, set target driver
        if cmds.objExists('{}.{}'.format(self.blendshape, target)):
            # Query the connection
            connections = cmds.listConnections(
                '{}.{}'.format(self.blendshape, target),
                plugs=True,
                source=True,
                destination=False
            )
            if connections:
                raise RuntimeError(
                    "{}'s target '{}' already has input connection!"
                    "Choose another target, or delete connection first!".format(
                        self.blendshape, target
                    )
                )

            # Math expression connect to blendshape target
            mathNodeCalculate.node_calculation(
                expression='{} + 0'.format(expression), affect_attr='{}.{}'.format(self.blendshape, target)
            )
            self.targets_connections_dict[target] = expression
            attrUtils.add_string_info(
                information=self.targets_connections_dict, obj=self.blendshape, attr='Targets_Connections_Dict'
            )

        else:
            raise RuntimeError('{} target does not exist in this blendshape node!'.format(target))

    def import_targets_connections(self, targets):
        """Import blendshape targets connections of given targets list,
           or import the whole existing targets connections.

        Args:
            targets (list/None): The list of targets to import targets connections,
                                 import all the targets connections if targets==None.

        """

        input_data = self.import_data()

        if input_data and 'targets_connections_dict' in input_data:

            if input_data['targets_connections_dict']:
                if targets:
                    targets_list = utils.list_operation(
                        list_a=targets, list_b=input_data['targets_connections_dict'].keys(), operation='&'
                    )
                    self.disconnect_targets_connections(targets=targets_list)
                    for target in targets_list:
                        try:
                            self.set_target_connection(
                                expression=input_data['targets_connections_dict'][target], target=target
                            )

                        except:
                            cmds.warning(
                                "Connection wasn't successfully set! "
                                "Please check if input objects in the expression exist in the scene."
                            )
                else:
                    self.set_targets_connections(input_connections=input_data['targets_connections_dict'])

                print("{}'s targets connections imported from {}".format(self.blendshape, current_path))

            else:
                cmds.warning("{}'s targets connections data doesn't exists".format(self.blendshape))

    def import_data(self):
        """Import blendshape data from json file, including
           'name', geometry', 'shapeType', 'members_indices',
           'targets_index_dict', 'targets_weights_dict', 'targets_connections_dict', 'targets_groups_list'.

        Returns:
            dict: Blendshape data, including
                  'name', geometry', 'shapeType', 'members_indices',
                  'targets_index_dict', 'targets_weights_dict',
                  'targets_connections_dict', 'targets_groups_list' as keys.
        """

        input_data = fileManage.read_file_info_from_json(path=current_path, user_file=self.blendshape_data_file)
        if input_data:
            print('{} imported from {}'.format(self.blendshape_data_file, current_path))
            return input_data

        else:
            cmds.warning("{} data {} file doesn't exists!".format(self.blendshape, self.blendshape_data_file))
            return {}


    def disconnect_targets_connections(self, targets):
        """Disconnect blendshape targets channels with any input connections

        Args:
            targets (list/None): Specific targets to get channels disconnected from input connection
                                 or all the targets to get channels disconnected from input connection if None.

        """

        targets_names = self.targets_connections_dict.keys()

        if targets:
            targets_names = utils.list_operation(
                list_a=targets, list_b=self.targets_connections_dict.keys(), operation='&'
            )

        for target_name in targets_names:
            # Query the connection
            connections = cmds.listConnections(
                '{}.{}'.format(self.blendshape, target_name),
                plugs=True,
                source=True,
                destination=False
            )
            if connections:
                cmds.disconnectAttr(connections[0], '{}.{}'.format(self.blendshape, target_name))

    def set_targets_connections(self, input_connections):
        """Connect blendshape targets channels with the expressions results outputs
           by using targets input connections dictionary.

        Args:
            input_connections (dict): Target input connections dictionary, with key(target name), value('expression').
                                      expression example: '2*(NodeA.attr+NodeB.attr)*NodeC.attr+NodeD.attr/NodeE.attr'.

        """

        for target, expression in input_connections.items():
            if cmds.objExists('{}.{}'.format(self.blendshape, target)):
                # Query the connection
                connections = cmds.listConnections(
                    '{}.{}'.format(self.blendshape, target),
                    plugs=True,
                    source=True,
                    destination=False
                )
                # Delete existing connection, refresh the whole thing by input data
                if connections:
                    cmds.disconnectAttr(connections[0], '{}.{}'.format(self.blendshape, target))

                # Try math expression connect to blendshape target
                try:
                    mathNodeCalculate.node_calculation(
                        expression='{} + 0'.format(expression), affect_attr='{}.{}'.format(self.blendshape, target)
                    )
                    self.targets_connections_dict[target] = expression

                except:
                    cmds.warning(
                        "Connection wasn't successfully set! "
                        "Please check if input objects in the expression exist in the scene."
                    )
                    continue

            else:
                cmds.warning(
                    "{} target does not exist in this blendshape node, connection process skipped.".format(target)
                )

        attrUtils.add_string_info(
            information=self.targets_connections_dict, obj=self.blendshape, attr='Targets_Connections_Dict'
        )
    def reconnect_targets_connections(self, targets):
        """Re-connect blendshape targets channels with
           output of expressions by using pre saved connection information.

        Args:
            targets (list/None): Specific targets to get channels reconnected from input connection
                                 or all the targets to get channels reconnected from input connection if None.

        """

        targets_names = self.targets_connections_dict.keys()

        if targets:
            targets_names = utils.list_operation(
                list_a=targets, list_b=self.targets_connections_dict.keys(), operation='&'
            )

        for target_name in targets_names:
            if cmds.objExists('{}.{}'.format(self.blendshape, target_name)):
                connections = cmds.listConnections(
                    '{}.{}'.format(self.blendshape, target_name),
                    plugs=True,
                    source=True,
                    destination=False
                )
                if connections:
                    cmds.disconnectAttr(connections[0], '{}.{}'.format(self.blendshape, target_name))

                # Try expression connect to blendshape target channel
                try:
                    mathNodeCalculate.node_calculation(
                        expression='{} + 0'.format(self.targets_connections_dict[target_name]),
                        affect_attr='{}.{}'.format(self.blendshape, target_name)
                    )

                except:
                    cmds.warning(
                        "target's connection wasn't successfully set, delete it from the connections dict."
                    )
                    del self.targets_connections_dict[target_name]

            else:
                cmds.warning(
                    "{} target no longer exists in this blendshape node, delete it from the connections dict.".format(
                        target_name
                    )
                )
                del self.targets_connections_dict[target_name]

        attrUtils.add_string_info(
            information=self.targets_connections_dict, obj=self.blendshape, attr='Targets_Connections_Dict'
        )


    def disconnect_targets_connections(self, targets):
        """Disconnect blendshape targets channels with any input connections

        Args:
            targets (list/None): Specific targets to get channels disconnected from input connection
                                 or all the targets to get channels disconnected from input connection if None.

        """

        targets_names = self.targets_connections_dict.keys()

        if targets:
            targets_names = utils.list_operation(
                list_a=targets, list_b=self.targets_connections_dict.keys(), operation='&'
            )

        for target_name in targets_names:
            # Query the connection
            connections = cmds.listConnections(
                '{}.{}'.format(self.blendshape, target_name),
                plugs=True,
                source=True,
                destination=False
            )
            if connections:
                cmds.disconnectAttr(connections[0], '{}.{}'.format(self.blendshape, target_name))
